"""
核反应堆子通道分析程序的核心分析类
"""

import numpy as np
from typing import Dict, List, Optional, Tuple
import time
import os

from constants.constants import (
    N_CHANNELS, CHANNEL_LENGTH, CHANNEL_TYPES, CHANNEL_GEOMETRY,
    CHANNEL_CONNECTIONS, INLET_PRESSURE, INLET_TEMP, INLET_MASS_FLUX,
    DEFAULT_TIME_STEP, SAVE_INTERVAL
)
from utils.helper_functions import initialize_arrays, calculate_geometry_parameters
from numerical.numba_functions import solve_momentum_equation, solve_energy_equation
from physics.physical_calculations import (
    calculate_fluid_properties,
    calculate_heat_transfer_coefficients,
    calculate_pressure_drop
)

class SubchannelAnalysis:
    """子通道分析的主类，处理所有计算和数据管理"""
    
    def __init__(self, 
                 n_axial: int = 2000,
                 total_time: float = 5.0,
                 time_step: float = DEFAULT_TIME_STEP,
                 save_interval: int = SAVE_INTERVAL,
                 results_folder: str = 'results'):
        """
        初始化子通道分析器
        
        参数:
            n_axial: 轴向节点数
            total_time: 总计算时间(s)
            time_step: 时间步长(s)
            save_interval: 保存结果的时间步间隔
            results_folder: 结果保存文件夹
        """
        # 基本参数
        self.n_axial = n_axial
        self.total_time = total_time
        self.time_step = time_step
        self.save_interval = save_interval
        self.results_folder = results_folder
        
        # 网格参数
        self.dz = CHANNEL_LENGTH / (n_axial - 1)
        self.z = np.linspace(0, CHANNEL_LENGTH, n_axial)
        
        # 初始化数组
        self.initialize_arrays()
        
        # 计算几何参数
        self.calculate_geometry()
        
        # 初始化历史记录
        self.history = {
            'time': [],
            'pressure': [],
            'temperature': [],
            'mass_flux': [],
            'void_fraction': [],
            'convergence': []
        }
        
        # 运行状态
        self.current_time = 0.0
        self.current_step = 0
        self.is_converged = False
    
    def initialize_arrays(self):
        """初始化所有计算所需的数组"""
        arrays = initialize_arrays(self.n_axial, N_CHANNELS)
        
        # 分配数组到实例变量
        self.pressure = arrays['pressure']
        self.temperature = arrays['temperature']
        self.mass_flux = arrays['mass_flux']
        self.void_fraction = arrays['void_fraction']
        self.quality = arrays['quality']
        self.enthalpy = arrays['enthalpy']
        self.density = arrays['density']
        self.velocity = arrays['velocity']
        
        # 设置初始条件
        self.pressure.fill(INLET_PRESSURE)
        self.temperature.fill(INLET_TEMP)
        self.mass_flux.fill(INLET_MASS_FLUX)
        self.void_fraction.fill(0.0)
        self.quality.fill(0.0)
    
    def calculate_geometry(self):
        """计算子通道几何参数"""
        self.geometry = calculate_geometry_parameters(
            CHANNEL_TYPES,
            CHANNEL_GEOMETRY,
            CHANNEL_CONNECTIONS
        )
    
    def run_transient_simulation(self):
        """运行瞬态模拟"""
        print(f"开始瞬态计算 (总时间: {self.total_time}s, 时间步长: {self.time_step}s)")
        
        start_time = time.time()
        n_steps = int(self.total_time / self.time_step)
        
        for step in range(n_steps):
            self.current_step = step
            self.current_time = step * self.time_step
            
            # 求解动量方程
            self.solve_momentum()
            
            # 求解���量方程
            self.solve_energy()
            
            # 更新流体属性
            self.update_fluid_properties()
            
            # 保存历史数据
            if step % self.save_interval == 0:
                self.save_history()
                print(f"计算进度: {self.current_time:.3f}s / {self.total_time}s")
        
        end_time = time.time()
        print(f"计算完成，耗时: {end_time - start_time:.2f}s")
    
    def solve_momentum(self):
        """求解动量方程"""
        self.pressure, self.mass_flux = solve_momentum_equation(
            self.pressure, self.mass_flux, self.density,
            self.geometry, self.dz, self.time_step
        )
    
    def solve_energy(self):
        """求解能量方程"""
        self.temperature, self.enthalpy = solve_energy_equation(
            self.temperature, self.enthalpy, self.mass_flux,
            self.geometry, self.dz, self.time_step
        )
    
    def update_fluid_properties(self):
        """更新流体属性"""
        properties = calculate_fluid_properties(
            self.pressure, self.temperature, self.quality
        )
        self.density = properties['density']
        self.velocity = properties['velocity']
        self.void_fraction = properties['void_fraction']
    
    def save_history(self):
        """保存当前时间步的计算结果到历史记录"""
        self.history['time'].append(self.current_time)
        self.history['pressure'].append(self.pressure.copy())
        self.history['temperature'].append(self.temperature.copy())
        self.history['mass_flux'].append(self.mass_flux.copy())
        self.history['void_fraction'].append(self.void_fraction.copy())
    
    def save_state(self, filename: str):
        """保存当前状态到文件"""
        np.savez(
            filename,
            time=self.current_time,
            step=self.current_step,
            pressure=self.pressure,
            temperature=self.temperature,
            mass_flux=self.mass_flux,
            void_fraction=self.void_fraction,
            quality=self.quality,
            enthalpy=self.enthalpy,
            density=self.density,
            velocity=self.velocity,
            history=self.history
        )
    
    def load_state(self, filename: str):
        """从文件加载状态"""
        data = np.load(filename)
        self.current_time = float(data['time'])
        self.current_step = int(data['step'])
        self.pressure = data['pressure']
        self.temperature = data['temperature']
        self.mass_flux = data['mass_flux']
        self.void_fraction = data['void_fraction']
        self.quality = data['quality']
        self.enthalpy = data['enthalpy']
        self.density = data['density']
        self.velocity = data['velocity']
        self.history = data['history'].item()
